'''
Author: duliang thinktanker@163.com
Date: 2025-08-30 22:11:35
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-08-31 13:35:59
FilePath: \watergrassyolo\rtsp_server.py
Description: RTSP视频转发服务器，支持多用户同时观看
'''
import cv2
import threading
import time
from flask import Flask, Response


class RTSPServer:

    def __init__(self, rtsp_url):
        self.rtsp_url = rtsp_url
        self.frame = None
        self.lock = threading.Lock()
        self.clients = []
        self.capture_thread = None
        self.running = False

    def start_capture(self):
        """启动视频捕获线程"""
        self.running = True
        self.capture_thread = threading.Thread(target=self._capture_frames)
        self.capture_thread.daemon = True
        self.capture_thread.start()

    def _capture_frames(self):
        """捕获RTSP视频帧"""
        cap = cv2.VideoCapture(self.rtsp_url)

        while self.running:
            ret, frame = cap.read()
            if ret:
                with self.lock:
                    self.frame = frame.copy()
            else:
                time.sleep(0.1)  # 如果读取失败，等待一段时间再重试

        cap.release()

    def stop_capture(self):
        """停止视频捕获"""
        self.running = False
        if self.capture_thread:
            self.capture_thread.join()

    def generate_frames(self):
        """生成视频帧用于HTTP流传输"""
        while self.running:
            with self.lock:
                if self.frame is not None:
                    ret, buffer = cv2.imencode('.jpg', self.frame)
                    if ret:
                        frame_bytes = buffer.tobytes()
                        yield (b'--frame\r\n'
                               b'Content-Type: image/jpg\r\n\r\n' +
                               frame_bytes + b'\r\n')
            time.sleep(0.03)  # 控制帧率约30fps


# 初始化Flask应用
app = Flask(__name__)

# 配置RTSP源URL
RTSP_URL = "rtsp://77.110.228.219/axis-media/media.amp"  # 请替换为实际的RTSP地址
server = RTSPServer(RTSP_URL)


@app.route('/')
def index():
    """主页面，显示视频流"""
    return '''
    <html>
        <head>
            <title>RTSP视频转发服务器</title>
        </head>
        <body>
            <h1>RTSP视频转发服务器</h1>
            <img src="/video_feed" width="800" height="600" />
        </body>
    </html>
    '''


@app.route('/video_feed')
def video_feed():
    """视频流端点"""
    return Response(server.generate_frames(),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


if __name__ == '__main__':
    # 启动视频捕获
    server.start_capture()

    # 启动Flask服务器
    # 0.0.0.0表示监听所有网络接口，便于多用户访问
    app.run(host='0.0.0.0', port=5000, threaded=True)
